# -*- coding: UTF-8 -*-
from __future__ import absolute_import, print_function

import tvm
from tvm import te

tgt_host = "c"
tgt = "c"

n = te.var("n")
A = te.placeholder((n,), name='A')
B = te.placeholder((n,), name='B')
C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")
print(type(C))

######################################################################
#   for (int i = 0; i < n; ++i) {
#     C[i] = A[i] + B[i];
#   }
s = te.create_schedule(C.op)
######################################################################
#   for (int bx = 0; bx < ceil(n / 64); ++bx) {
#     for (int tx = 0; tx < 64; ++tx) {
#       int i = bx * 64 + tx;
#       if (i < n) {
#         C[i] = A[i] + B[i];
#       }
#     }
#   }
######################################################################
bx, tx = s[C].split(C.op.axis[0], factor=4)
######################################################################
if tgt == "cuda" or tgt == "rocm" or tgt.startswith('opencl'):
  s[C].bind(bx, te.thread_axis("blockIdx.x"))
  s[C].bind(tx, te.thread_axis("threadIdx.x"))
######################################################################
fadd = tvm.build(s, [A, B, C], tgt, target_host=tgt_host, name="myadd")
print(fadd.get_source())
